import sys
from PyQt6.QtWidgets import QApplication, QWidget, QVBoxLayout, QProgressBar, QPushButton
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.linear_model import Ridge
from sklearn.svm import SVR
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from concurrent.futures import ProcessPoolExecutor, as_completed
import time

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Model Training Progress")
        self.setGeometry(100, 100, 300, 200)

        self.layout = QVBoxLayout(self)

        self.progressBarAll = QProgressBar(self)
        self.progressBarSub = QProgressBar(self)
        self.layout.addWidget(self.progressBarAll)
        self.layout.addWidget(self.progressBarSub)

        self.startButton = QPushButton("Start Training", self)
        self.startButton.clicked.connect(self.start_training)
        self.layout.addWidget(self.startButton)

        self.setLayout(self.layout)

    def update_progress(self, progress):
        self.progressBarSub.setValue(progress)

    def update_progress_all(self, completed_tasks, total_tasks):
        self.progressBarAll.setValue(int((completed_tasks / total_tasks) * 100))

    def train_model(self, model_class, *args):
        model = model_class(*args)
        for i in range(1, 101):  # 模拟训练过程，100个步骤
            time.sleep(0.05)  # 模拟每一步的训练时间
            self.update_progress(i)  # 更新分进度条
        return model

    def start_training(self):
        self.progressBarAll.setValue(0)
        self.progressBarSub.setValue(0)

        # 创建样本数据集
        X, y = make_regression(n_samples=1000, n_features=20, noise=0.1, random_state=42)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

        models = [
            (GradientBoostingRegressor, {}),
            (Ridge, {"alpha": 1.0}),
            (SVR, {"C": 1.0, "kernel": "rbf"})
        ]

        total_models = len(models)

        with ProcessPoolExecutor() as executor:
            futures = {executor.submit(self.train_model, model_class, **params): index for index, (model_class, params) in enumerate(models)}

            for future in as_completed(futures):
                index = futures[future]
                self.update_progress_all(index + 1, total_models)  # 更新总进度条

                try:
                    model = future.result()
                    print(f"Model {index + 1} trained: {model}")
                except Exception as exc:
                    print(f"Model {index + 1} generated an exception: {exc}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
